#!/usr/bin/env python
# -*- coding: utf-8 -*-

import abc
import logging
import hashlib
import time
from datetime import datetime
from typing import Dict, List, Any, Tuple, Optional, Union
import re

# 获取logger而不是配置它
logger = logging.getLogger(__name__)

class GraphConnection(abc.ABC):
    """
    图数据库连接的抽象基类，定义图数据库操作的通用接口
    """
    
    @abc.abstractmethod
    def connect(self) -> bool:
        """
        连接到图数据库
        
        Returns:
            bool: 连接是否成功
        """
        pass
    
    @abc.abstractmethod
    def close(self) -> None:
        """
        关闭图数据库连接
        """
        pass
    
    @abc.abstractmethod
    def execute_query(self, query: str) -> Tuple[bool, Any]:
        """
        执行查询语句
        
        Args:
            query (str): 查询语句
            
        Returns:
            Tuple[bool, Any]: (是否成功, 查询结果或错误信息)
        """
        pass
    
    @abc.abstractmethod
    def create_vertex(self, label: str, vertex_id: str, properties: Dict[str, Any]) -> bool:
        """
        创建顶点
        
        Args:
            label (str): 顶点标签/类型
            vertex_id (str): 顶点ID
            properties (Dict[str, Any]): 顶点属性
            
        Returns:
            bool: 操作是否成功
        """
        pass
    
    @abc.abstractmethod
    def create_edge(self, label: str, from_id: str, to_id: str, properties: Dict[str, Any]) -> bool:
        """
        创建边
        
        Args:
            label (str): 边标签/类型
            from_id (str): 起始顶点ID
            to_id (str): 终止顶点ID
            properties (Dict[str, Any]): 边属性
            
        Returns:
            bool: 操作是否成功
        """
        pass
    
    @abc.abstractmethod
    def use_graph(self, graph_name: str) -> bool:
        """
        选择要使用的图
        
        Args:
            graph_name (str): 图名称
            
        Returns:
            bool: 操作是否成功
        """
        pass


class NebulaGraphConnection(GraphConnection):
    """
    NebulaGraph数据库连接实现
    """
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化NebulaGraph连接
        
        Args:
            config (Dict[str, Any]): 连接配置，包含以下字段:
                - ip: 服务器IP地址
                - port: 服务器端口
                - user: 用户名
                - password: 密码
                - space: 工作空间名称
        """
        self.config = config
        self.client = None
        self.pool = None
        self.vid_cache = {}  # 缓存已经生成的VID，避免重复计算
    
    def connect(self) -> bool:
        """
        连接到NebulaGraph数据库
        
        Returns:
            bool: 连接是否成功
        """
        try:
            from nebula3.gclient.net import ConnectionPool
            from nebula3.Config import Config
            
            # 创建连接配置
            nebula_config = Config()
            nebula_config.max_connection_pool_size = 10
            
            # 创建连接池
            connection_pool = ConnectionPool()
            ok = connection_pool.init([
                (self.config["ip"], self.config["port"])
            ], nebula_config)
            
            if not ok:
                logger.error("初始化NebulaGraph连接池失败")
                return False
                
            # 获取会话
            client = connection_pool.get_session(
                self.config["user"], 
                self.config["password"]
            )
            
            # 如果指定了工作空间，则选择它
            if "space" in self.config and self.config["space"]:
                resp = client.execute(f"USE {self.config['space']}")
                if not resp.is_succeeded():
                    logger.error(f"选择工作空间失败: {resp.error_msg()}")
                    client.release()
                    connection_pool.close()
                    return False
            
            self.client = client
            self.pool = connection_pool
            logger.info(f"已连接到NebulaGraph数据库, 工作空间: {self.config.get('space', 'default')}")
            return True
        except Exception as e:
            logger.error(f"连接NebulaGraph失败: {str(e)}")
            return False
    
    def close(self) -> None:
        """
        关闭NebulaGraph连接
        """
        try:
            if self.client:
                self.client.release()
            if self.pool:
                self.pool.close()
            logger.info("已关闭NebulaGraph连接")
        except Exception as e:
            logger.error(f"关闭NebulaGraph连接失败: {str(e)}")
    
    def execute_query(self, query: str) -> Tuple[bool, Any]:
        """
        执行NebulaGraph查询语句
        
        Args:
            query (str): nGQL查询语句
            
        Returns:
            Tuple[bool, Any]: (是否成功, 查询结果或错误信息)
        """
        try:
            if not self.client:
                return False, "未连接到NebulaGraph数据库"
            
            # 记录要执行的查询（但截断过长的查询）    
            log_query = query if len(query) < 1000 else query[:1000] + "...(截断)"
            logger.debug(f"执行查询: {log_query}")
                
            resp = self.client.execute(query)
            if resp.is_succeeded():
                logger.debug(f"查询成功执行")
                return True, resp
            else:
                error_msg = resp.error_msg()
                logger.error(f"查询执行失败: {error_msg}")
                return False, error_msg
        except Exception as e:
            logger.error(f"执行NebulaGraph查询失败: {str(e)}")
            return False, str(e)
    
    def create_vertex(self, label: str, vertex_id: str, properties: Dict[str, Any]) -> bool:
        """
        在NebulaGraph中创建顶点
        
        Args:
            label (str): 顶点标签/类型
            vertex_id (str): 顶点ID (会被转换为32位字符串格式)
            properties (Dict[str, Any]): 顶点属性
            
        Returns:
            bool: 操作是否成功
        """
        try:
            # 生成32位VID
            nebula_vid = self._generate_vid(vertex_id)
            logger.debug(f"原始ID: {vertex_id} -> NebulaGraph VID: {nebula_vid}")
            
            # 转换timestamp格式字段
            self._convert_timestamp_fields(properties)
            
            # 过滤掉值为None的属性
            filtered_properties = {k: v for k, v in properties.items() if v is not None}
            
            if not filtered_properties:
                # 如果没有属性，使用简化语法
                query = f'INSERT VERTEX {label}() VALUES "{nebula_vid}":() '
            else:
                # 构建属性列表和值列表
                prop_names = ', '.join(filtered_properties.keys())
                prop_values = ', '.join([self._format_value(v) for v in filtered_properties.values()])
                
                # 构建INSERT VERTEX语句
                query = f'INSERT VERTEX {label} ({prop_names}) VALUES "{nebula_vid}":({prop_values})'
            
            success, result = self.execute_query(query)
            if not success:
                logger.warning(f"创建顶点 {label}:{vertex_id} 失败: {result}")
            return success
        except Exception as e:
            logger.error(f"创建顶点 {label}:{vertex_id} 失败，异常: {str(e)}")
            return False
    
    def create_edge(self, label: str, from_id: str, to_id: str, properties: Dict[str, Any]) -> bool:
        """
        在NebulaGraph中创建边
        
        Args:
            label (str): 边标签/类型
            from_id (str): 起始顶点ID (会被转换为32位字符串格式)
            to_id (str): 终止顶点ID (会被转换为32位字符串格式)
            properties (Dict[str, Any]): 边属性
            
        Returns:
            bool: 操作是否成功
        """
        try:
            # 生成32位VID
            from_vid = self._generate_vid(from_id)
            to_vid = self._generate_vid(to_id)
            
            # 记录原始ID与VID的映射关系，便于调试
            logger.debug(f"边关系ID转换: {from_id} -> {from_vid}, {to_id} -> {to_vid}")
            
            # 转换timestamp格式字段
            self._convert_timestamp_fields(properties)
            
            # 过滤掉值为None的属性
            filtered_properties = {k: v for k, v in properties.items() if v is not None}
            
            if not filtered_properties:
                # 如果没有属性，使用简化语法
                query = f'INSERT EDGE {label}() VALUES "{from_vid}" -> "{to_vid}":() '
            else:
                # 构建属性列表和值列表
                prop_names = ', '.join(filtered_properties.keys())
                prop_values = ', '.join([self._format_value(v) for v in filtered_properties.values()])
                
                # 构建INSERT EDGE语句
                query = f'INSERT EDGE {label} ({prop_names}) VALUES "{from_vid}" -> "{to_vid}":({prop_values})'
            
            success, result = self.execute_query(query)
            if not success:
                logger.warning(f"创建边 {label}: {from_id} -> {to_id} 失败: {result}")
            else:
                logger.debug(f"成功创建边 {label}: {from_id} -> {to_id}")
            return success
        except Exception as e:
            logger.error(f"创建边 {label}: {from_id} -> {to_id} 失败，异常: {str(e)}")
            return False
    
    def use_graph(self, graph_name: str) -> bool:
        """
        选择要使用的NebulaGraph空间
        
        Args:
            graph_name (str): 空间名称
            
        Returns:
            bool: 操作是否成功
        """
        try:
            if not self.client:
                return False
                
            resp = self.client.execute(f"USE {graph_name}")
            if resp.is_succeeded():
                self.config["space"] = graph_name
                return True
            else:
                logger.error(f"选择工作空间失败: {resp.error_msg()}")
                return False
        except Exception as e:
            logger.error(f"选择工作空间失败: {str(e)}")
            return False
    
    def _generate_vid(self, original_id: str) -> str:
        """
        生成NebulaGraph需要的32位VID字符串
        
        Args:
            original_id (str): 原始ID字符串
            
        Returns:
            str: 32位固定长度的字符串
        """
        # 使用缓存避免重复计算相同ID
        if original_id in self.vid_cache:
            return self.vid_cache[original_id]
        
        # 使用MD5生成32位字符串，确保符合NebulaGraph VID长度要求
        vid = hashlib.md5(original_id.encode('utf-8')).hexdigest()
        
        # 记录到缓存
        self.vid_cache[original_id] = vid
        return vid
    
    def _convert_timestamp_fields(self, properties: Dict[str, Any]) -> None:
        """
        将属性中的时间字段转换为timestamp格式
        
        Args:
            properties: 属性字典，会被原地修改
        """
        # 需要转换为timestamp的字段列表
        timestamp_fields = [
            "last_updated", "last_modified", "created_at", "updated_at", 
            "commit_date", "timestamp"
        ]
        
        for field in timestamp_fields:
            if field in properties:
                value = properties[field]
                if isinstance(value, datetime):
                    # 转换为Unix时间戳(秒)
                    properties[field] = int(value.timestamp())
                elif isinstance(value, str):
                    try:
                        # 尝试将字符串解析为日期时间，然后转换为时间戳
                        dt = datetime.fromisoformat(value.replace('Z', '+00:00'))
                        properties[field] = int(dt.timestamp())
                    except (ValueError, TypeError):
                        # 如果无法解析，保留原值
                        pass
    
    def _format_value(self, value: Any) -> str:
        """
        格式化属性值为NebulaGraph支持的字符串格式，并处理编码问题
        
        Args:
            value (Any): 属性值
            
        Returns:
            str: 格式化后的字符串
        """
        if value is None:
            return "NULL"
        elif isinstance(value, bool):
            return str(value).lower()
        elif isinstance(value, (int, float)):
            return str(value)
        elif isinstance(value, str):
            # 解决编码问题和多行文本问题
            try:
                # 1. 移除字符串中的换行符和多余空格
                cleaned_value = value.replace('\n', ' ').replace('\r', ' ')
                
                # 2. 将连续的多个空格替换为单个空格
                cleaned_value = re.sub(r'\s+', ' ', cleaned_value)
                
                # 3. 确保字符串是有效的UTF-8
                cleaned_value = cleaned_value.encode('utf-8', errors='ignore').decode('utf-8')
                
                # 4. 转义特殊字符，特别是引号
                escaped = cleaned_value.replace('"', '\\"').replace("'", "\\'")
                
                return f'"{escaped}"'
            except Exception as e:
                logger.warning(f"字符串清理过程中出错: {str(e)}")
                # 返回一个安全值
                return '""'
        elif isinstance(value, datetime):
            # 转换为时间戳(秒)
            return str(int(value.timestamp()))
        elif isinstance(value, (list, dict)):
            import json
            try:
                # 将复杂对象转为JSON字符串，并确保双引号被正确转义
                json_str = json.dumps(value)
                # 处理JSON字符串
                cleaned_json = json_str.encode('utf-8', errors='ignore').decode('utf-8')
                escaped_json = cleaned_json.replace('"', '\\"')
                return f'"{escaped_json}"'
            except Exception as e:
                logger.warning(f"JSON格式化过程中出错: {str(e)}")
                return '""'
        else:
            # 其他类型转为字符串
            try:
                str_value = str(value)
                # 移除换行符和处理编码
                str_value = str_value.replace('\n', ' ').replace('\r', ' ')
                str_value = re.sub(r'\s+', ' ', str_value)
                cleaned_value = str_value.encode('utf-8', errors='ignore').decode('utf-8')
                return f'"{cleaned_value}"'
            except Exception as e:
                logger.warning(f"值转换为字符串时出错: {str(e)}")
                return '""'


class Neo4jConnection(GraphConnection):
    """
    Neo4j数据库连接实现 (未完成实现)
    """
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化Neo4j连接
        
        Args:
            config (Dict[str, Any]): 连接配置
        """
        self.config = config
        # 这里只是一个示例，实际实现时需要添加Neo4j连接的相关字段
    
    def connect(self) -> bool:
        """
        连接到Neo4j数据库 (未实现)
        
        Returns:
            bool: 连接是否成功
        """
        # 未实现
        logger.warning("Neo4j连接尚未实现")
        return False
    
    def close(self) -> None:
        """
        关闭Neo4j连接 (未实现)
        """
        # 未实现
        pass
    
    def execute_query(self, query: str) -> Tuple[bool, Any]:
        """
        执行Neo4j查询语句 (未实现)
        
        Args:
            query (str): Cypher查询语句
            
        Returns:
            Tuple[bool, Any]: (是否成功, 查询结果或错误信息)
        """
        # 未实现
        return False, "Neo4j连接方法尚未实现"
    
    def create_vertex(self, label: str, vertex_id: str, properties: Dict[str, Any]) -> bool:
        """
        在Neo4j中创建顶点 (未实现)
        
        Args:
            label (str): 顶点标签/类型
            vertex_id (str): 顶点ID
            properties (Dict[str, Any]): 顶点属性
            
        Returns:
            bool: 操作是否成功
        """
        # 未实现
        return False
    
    def create_edge(self, label: str, from_id: str, to_id: str, properties: Dict[str, Any]) -> bool:
        """
        在Neo4j中创建边 (未实现)
        
        Args:
            label (str): 边标签/类型
            from_id (str): 起始顶点ID
            to_id (str): 终止顶点ID
            properties (Dict[str, Any]): 边属性
            
        Returns:
            bool: 操作是否成功
        """
        # 未实现
        return False
    
    def use_graph(self, graph_name: str) -> bool:
        """
        选择要使用的Neo4j数据库 (未实现)
        
        Args:
            graph_name (str): 数据库名称
            
        Returns:
            bool: 操作是否成功
        """
        # 未实现
        return False


def create_connection(db_type: str, config: Dict[str, Any]) -> GraphConnection:
    """
    创建图数据库连接
    
    Args:
        db_type (str): 数据库类型，如"nebula"或"neo4j"
        config (Dict[str, Any]): 连接配置
        
    Returns:
        GraphConnection: 图数据库连接实例
        
    Raises:
        ValueError: 如果指定的数据库类型不支持
    """
    if db_type.lower() == "nebula":
        return NebulaGraphConnection(config)
    elif db_type.lower() == "neo4j":
        return Neo4jConnection(config)
    else:
        raise ValueError(f"不支持的数据库类型: {db_type}") 